import os
import cv2
from sklearn import preprocessing
from sklearn import decomposition
import numpy as np
import time

print("Getting trainning photos information...")
img_path = []
labels = []

dpath = './train'

folders = os.listdir(dpath)
for f in folders:
    imgs = os.listdir(dpath + "/" + f)
    for i in imgs:
        img_path.append(dpath + "/" + f + "/" + i)
        labels.append(f)

plabel = preprocessing.LabelEncoder()
y_label = plabel.fit_transform(labels)

train_img = []
train_labels = []

print("Detecting photos face features...")
n = 0
classifier = cv2.CascadeClassifier("./haarcascade_frontalface_alt.xml")
for i in img_path:
    img = cv2.imread(i)
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    face_area = classifier.detectMultiScale(img_gray, 1.5, 5)
    for (x, y, w, h) in face_area:
        face = img_gray[y: y + h, x: x + w]
        train_img.append(face)
        train_labels.append(y_label[n])
    n+=1

train_labels = np.array(train_labels)

print("PCA handling photos's features.......")
#PCA 降维
train_img2 = []
for i in range(0, len(train_img)):
    pca = decomposition.PCA()
    pca.fit(train_img[i])
    ft1 = pca.explained_variance_
    #ft2_num = len(np.where(ft1 > 0.6)[0])
    #print(str(ft2_num))
    pca.n_components = 67
    this_img = pca.fit_transform(train_img[i])
    train_img2.append(this_img)

print("Predicting photos's labels...")
#LBP, eigenface, fisherface, currently using LBP

recongnizer = cv2.face.LBPHFaceRecognizer_create()
recongnizer.train(train_img2, train_labels)

'''
test_path = "./test"
all_test = os.listdir(test_path)

for i in all_test:
    img_path = test_path + "/" + i
    img = cv2.imread(img_path)
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    face_area = classifier.detectMultiScale(img_gray, 1.5, 5)
    for (x, y, w, h) in face_area:
        rst, b = recongnizer.predict(img_gray[y:y+h, x:x+w])

    rst2 = plabel.inverse_transform([rst])[0]
    print("This person is " + str(rst2))
    cv2.putText(img, str(rst2), (20,20), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255), 1, False)
    cv2.imshow(str(rst2), img)
    cv2.waitKey(1)
    time.sleep(3)
'''

cap = cv2.VideoCapture(0)
while True:
    ret, img = cap.read()
    img = cv2.resize(img, None, fx = 1, fy = 1)
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    face_area = classifier.detectMultiScale(img_gray, 1.5, 5)
    for (x, y, w, h) in face_area:
        rst, b = recongnizer.predict(img_gray[y:y+h, x:x+w])
        rst2 = plabel.inverse_transform([rst])[0]
        cv2.putText(img, str(rst2), (20,20), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255), 1, False)
        
    cv2.imshow("img", img)
    time.sleep(0.1)
    key = cv2.waitKey(1)
    if(key == ord('q')):
        print("Force cancel!")
        break
    
print("Quit program and close all windows...")
cap.release()
cv2.destroyAllWindows()
